import random
from parameter import *
import numpy as np
import os
import networkx as nx

exit_list = []
target_list = []
robot_list = []

for index in range(152):
    adj_dir = train_adj_path
    adj_list = os.listdir(adj_dir)
    adj_list.sort(reverse=False)
    adj_index = index % np.size(adj_list)
    file_path = adj_dir + '/' + adj_list[adj_index]

    with open(file_path, 'r') as f:
        node_num = int(f.readline().strip())
        adj_matrix = []
        for line in f:
            adj_matrix.append(list(map(int, line.split())))
        adj_matrix = np.array(adj_matrix)
    graph = nx.from_numpy_array(adj_matrix)
    network_adjacent_matrix = nx.floyd_warshall_numpy(graph)

    next_node = [[j for j in range(node_num)] for i in range(node_num)]
    for i in range(node_num):
        for j in range(node_num):
            for k in range(node_num):
                if network_adjacent_matrix[i][k] == 1 and network_adjacent_matrix[i][j] == network_adjacent_matrix[k][j] + 1:
                    next_node[i][j] = k
    next_node = np.array(next_node)

    if EXIT_NUM > 0:       
        while True:
            # 随机选择出口
            exit_indexes = random.sample(range(node_num), EXIT_NUM)

            # 选择追捕者位置，排除出口节点
            candidates = [index for index in range(node_num) if index not in exit_indexes]
            robot_indexes = random.sample(candidates, N_ROBOTS)

            # 固定逃逸者到最近出口的距离为指定值
            min_dist = MIN_EVADOR_EXIT_DIST  # 设置逃逸者到最近出口的固定距离

            # 筛选符合条件的逃逸者节点
            valid_target_candidates = [
                node for node in candidates
                if any(
                    network_adjacent_matrix[node][exit_index] == min_dist
                    for exit_index in exit_indexes
                )
            ]

            if not valid_target_candidates:
                print("No valid target nodes found for the desired distance. Adjust parameters.")
                continue  # 跳过当前循环，重新选择出口和追捕者

            # 随机选择一个目标节点
            target_index = random.choice(valid_target_candidates)

            # 找出逃逸者最近的出口
            closest_exit_index = min(
                exit_indexes,
                key=lambda exit_index: network_adjacent_matrix[target_index][exit_index]
            )
            closest_exit_distance = network_adjacent_matrix[target_index][closest_exit_index]

            # 检查条件：最近出口距离必须等于 DESIRED_DISTANCE
            if closest_exit_distance != min_dist:
                continue

            # # 检查条件：最近出口是否有追捕者比逃逸者更接近
            # closer_robots = [
            #     robot_index for robot_index in robot_indexes
            #     if self.network_adjacent_matrix[robot_index][closest_exit_index] < closest_exit_distance
            # ]

            # if not closer_robots:  # 如果没有追捕者比逃逸者更接近最近出口
            #     continue


            # 检查条件 2：每个距离逃逸者小于 10 的出口，都有追捕者更近
            valid_exits = [
                exit_index for exit_index in exit_indexes
                if network_adjacent_matrix[target_index][exit_index] <= 10
            ]

            all_exits_satisfied = True
            for exit_index in valid_exits:
                closer_robots = [
                    robot_index for robot_index in robot_indexes
                    if network_adjacent_matrix[robot_index][exit_index] < network_adjacent_matrix[target_index][exit_index]
                ]
                if not closer_robots:
                    all_exits_satisfied = False
                    break

            if not all_exits_satisfied:
                continue

            # 如果满足所有条件，退出循环
            break
            
    else:
        raise ValueError("Must Have Exit")
    
    # print('file', file_path, 'exit: ', exit_indexes, 'pursuer: ', robot_indexes, 'evader: ', target_index)
    exit_list.append(exit_indexes)
    target_list.append(target_index)
    robot_list.append(robot_indexes)
    # dist_list.append(network_adjacent_matrix)
    # next_list.append(next_node)
    file = 'init/train/'+ str(adj_index) +'.npz'
    np.savez(file, dist=network_adjacent_matrix, next=next_node)
    data = np.load(file)
    print(data['next'].shape)
    print(data['dist'].shape)

exit_list = np.array(exit_list)
target_list = np.array(target_list)
robot_list = np.array(robot_list)

np.savez('init/train.npz', exit=exit_list, target=target_list, robot=robot_list)

data = np.load('init/train.npz')
print(data['exit'].shape)
print(data['target'].shape)
print(data['robot'].shape)
